from contextlib import contextmanager
import logging
from typing import Generator, Optional
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.exc import SQLAlchemyError

from src.database.models import Base
from src.log.logger import setup_logger

class Database:
    """数据库管理类"""
    
    _instance = None
    
    def __init__(self, database_url: str):
        self.logger = setup_logger("database")
        self.engine = create_engine(database_url)
        self.SessionLocal = sessionmaker(bind=self.engine)
        
    @classmethod
    def initialize(cls, database_url: str) -> 'Database':
        """初始化数据库实例"""
        if cls._instance is None:
            cls._instance = cls(database_url)
            # 创建所有表
            Base.metadata.create_all(cls._instance.engine)
        return cls._instance
    
    @classmethod
    def get_instance(cls) -> 'Database':
        """获取数据库实例"""
        if cls._instance is None:
            raise RuntimeError("Database not initialized")
        return cls._instance
    
    @contextmanager
    def get_session(self) -> Generator[Session, None, None]:
        """获取数据库会话的上下文管理器"""
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except SQLAlchemyError as e:
            session.rollback()
            self.logger.error(f"数据库操作失败: {e}", exc_info=True)
            raise
        finally:
            session.close() 